#include <assert.h>
#include <omp.h>
#include <stdio.h>

#include <algorithm>

#include "oblivious.h"


using namespace std;

static const bool VERBOSE = false;

//static void mat_mult_rec(const MatrixSlice a, const MatrixSlice b, MatrixSlice c, int base_case_size);

/*
extern "C" void cblas_dgemm(CBLAS_ORDER, CBLAS_TRANSPOSE, CBLAS_TRANSPOSE, int, int, int, double, const double*, int, const double*, int, double, double*, int);
*/

Matrix::Matrix(int rows, int columns) {
  this->rows = rows;
  this->columns = columns;
  this->d = new double[rows * columns];
}

Matrix::Matrix(double* data, int rows, int columns) {
  this->rows = rows;
  this->columns = columns;
  this->d = data;
}

Matrix::Matrix() {
  rows = columns = 0;
  d = NULL;
}

//Matrix::~Matrix() {
  //delete [] d;
//}

void Matrix::resize(int new_rows, int new_columns) {
  this->rows = new_rows;
  this->columns = new_columns;
  this->d = new double[new_rows * new_columns];
}

void Matrix::zero() {
  for (int i = 0; i < rows * columns; i++)
    d[i] = 0.0;
}

MatrixSlice Matrix::slice() const { 
  return MatrixSlice(this); 
}

void Matrix::dump() const {
  for (int row = 0; row < rows; row++) {
    for (int col = 0; col < columns; col++) {
      if (row == 0 && col == 0)
        printf("[%f ", d[row * columns + col]);
      else
        printf(", %f ", d[row * columns + col]);
    }
    printf("\n");
  }
  printf("]\n");
}

void MatrixSlice::dump() const {
  
  for (int row_ind = 0; row_ind < rows; row_ind++) {
    for (int col_ind = 0; col_ind < columns; col_ind++) {
      if (row_ind == 0 && col_ind == 0)
        printf("[%f ", at(row_ind, col_ind));
      else
        printf(", %f ", at(row_ind, col_ind));
    } 
    printf("\n");
  }
} 

MatrixSlice::MatrixSlice() {
  this->parent = NULL;
  
  d = NULL;
  rows = 0;
  columns = 0;
  row_stride = 0;
}

MatrixSlice::MatrixSlice(const Matrix *parent) {
  this->parent = parent;

  d = parent->d;
  rows = parent->rows;
  columns = parent->columns;
  // this needs to be columns
  row_stride = parent->columns;
  
}

vector<MatrixSlice> MatrixSlice::split_rows() const {
  MatrixSlice result[2];
  
  /*
  for (int i = 0; i < 2; i++) {
    MatrixSlice *slice = result + i;
    
    slice->parent = parent;
    slice->d = d + i * rows * row_stride / 2;
    slice->rows = rows / 2;
    slice->columns = columns;
    slice->row_stride = row_stride;
  }
  */
  
  MatrixSlice *top_slice = result;
  top_slice->parent = parent;
  top_slice->d = d;
  top_slice->rows = rows / 2;
  top_slice->columns = columns;
  top_slice->row_stride = row_stride;
  
  MatrixSlice *bottom_slice = result +1;
  bottom_slice->parent = parent;
  bottom_slice->d = d + top_slice->rows * row_stride;
  bottom_slice->rows = rows - top_slice->rows;
  bottom_slice->columns = columns;
  bottom_slice->row_stride = row_stride;
  
  
  return vector<MatrixSlice>(result, result + 2);
}

vector<MatrixSlice> MatrixSlice::split_columns() const {
  MatrixSlice result[2];
  
  /*
  for (int i = 0; i < 2; i++) {
    MatrixSlice *slice = result + i;
    
    slice->parent = parent;
    slice->d = d + i * columns / 2 ;
    slice->rows = rows;
    slice->columns = columns / 2;
    slice->row_stride = row_stride;
  }
   */
  
  MatrixSlice *left_slice = result;
  left_slice->parent = parent;
  left_slice->d = d;
  left_slice->rows = rows;
  left_slice->columns = columns / 2;
  left_slice->row_stride = row_stride;
  
  MatrixSlice *right_slice = result +1;
  right_slice->parent = parent;
  right_slice->d = d + left_slice->columns;
  right_slice->rows = rows;
  right_slice->columns = columns - left_slice->columns;
  right_slice->row_stride = row_stride;
  
  return vector<MatrixSlice>(result, result + 2); 
}

void mat_mult_rec(const MatrixSlice a, const MatrixSlice b, MatrixSlice c, int base_case_size) {
  
  // printf("%d/%d\n", omp_get_thread_num(), omp_get_num_threads());
  
  if (VERBOSE)
    printf("mat_mult_rec: a: (%d, %d); b: (%d, %d); c: (%d, %d), t: %d\n", 
      a.get_rows(), a.get_columns(),
      b.get_rows(), b.get_columns(),
	   c.get_rows(), c.get_columns(), omp_get_thread_num());

  /*
    if (c.get_columns() == 0) {
	     printf("Empty c\n");
	   }
  */
  // assert(a.get_columns() == b.get_rows());
  // assert(a.get_rows() > 0 && a.get_columns() > 0);
  // assert(b.get_rows() > 0 && b.get_columns() > 0);
  // assert(c.get_rows() > 0 && c.get_columns() > 0);


  int m = a.get_rows();
  int n = a.get_columns();
  int p = b.get_columns();

  int max_size = max(m, max(n, p));
  
    //if (m == n == p == 1) {
    //if ((m == 1) && (n == 1) && (p == 1)) {
  if (max_size <= base_case_size) {
    //printf("Base Case, m = %d\n", m);

    // For BLAS, the matrices must live in contiguous memory

    double* a_copy = (double*)malloc(sizeof(double) * a.get_rows() * a.get_columns());
    double* b_copy = (double*)malloc(sizeof(double) * b.get_rows() * b.get_columns());
    double* c_copy = (double*)malloc(sizeof(double) * c.get_rows() * c.get_columns());
    
    for (int i = 0; i < a.get_rows(); i++) {
      std::memcpy(a_copy+(i*a.get_columns()), a.get_ptr() + (i*a.get_row_stride()), a.get_columns() * sizeof(double));
    }

    for (int i = 0; i < b.get_rows(); i++) {
      std::memcpy(b_copy+(i*b.get_columns()), b.get_ptr() + (i*b.get_row_stride()), b.get_columns() * sizeof(double));
    }
    

    //c.at(0, 0) += a.at(0, 0) * b.at(0, 0);
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, a.get_rows(), b.get_columns(), a.get_columns(), 1.0, a_copy, a.get_columns(), b_copy, b.get_columns(), 0.0, c_copy, c.get_columns());
    
    // now, add the new c into the old one
#pragma omp critical(dataupdate)
    {
      for (int i = 0; i < c.get_rows(); i++) {
	cblas_daxpy(c.get_columns(), 1.0, c_copy + (i*c.get_columns()), 1, c.get_ptr()+(i*c.get_row_stride()), 1);
      } 
    } // critical

    free(a_copy);
    free(b_copy);
    free(c_copy);
    
    } else if (m >= max(n, p)) {    

      // Split A along the row axis.
      // printf("Spitting A\n");
      vector<MatrixSlice> as = a.split_rows();
      vector<MatrixSlice> cs = c.split_rows();

#pragma omp task // firstprivate(as, cs)
      {
        // MatrixSlice b_copy(b);
        mat_mult_rec(as[0], b, cs[0], base_case_size);
      }
      
      // #pragma omp task 
      mat_mult_rec(as[1], b, cs[1], base_case_size);


    } else if (n >= max(m, p)) {
      // Split A along the column axis and B along the row axis.
      //printf("Splitting both\n");
      vector<MatrixSlice> as = a.split_columns();
      vector<MatrixSlice> bs = b.split_rows();


#pragma omp task // firstprivate(as, bs)
      {
        // MatrixSlice c_copy(c);
        mat_mult_rec(as[0], bs[0], c, base_case_size);
      }
      
      // #pragma omp task
      mat_mult_rec(as[1], bs[1], c, base_case_size);
     
    } else {
      // assert(p >= max(m, n));
      //printf("Splitting B\n");
      // Split B along the column axis.
      vector<MatrixSlice> bs = b.split_columns();
      vector<MatrixSlice> cs = c.split_columns();

#pragma omp task // firstprivate(bs, cs)
      {
        // MatrixSlice a_copy(a);
        mat_mult_rec(a, bs[0], cs[0], base_case_size);
      }
    
      // #pragma omp task
      mat_mult_rec(a, bs[1], cs[1], base_case_size);
    
    }

    #pragma omp taskwait
  
}

void mat_mult_rec(const Matrix &a, const Matrix &b, Matrix &c, int base_case_size) {
  
  // c.resize(a.get_rows(), b.get_columns());
  // c.zero(); // Important b/c we may add into c recursively.

#pragma omp parallel
  #pragma omp single
  {
    //#pragma omp task

    mat_mult_rec(a.slice(), b.slice(), c.slice(), base_case_size);
  }
}

